# -*-coding=utf-8 -*-
from sklearn.feature_extraction import DictVectorizer
import csv
from sklearn import preprocessing
from sklearn import tree
from sklearn.externals.six import StringIO

#Read in the csv file and put features in a list of dict and list of class label.
allElectronicsData = open(r'datasets/DecisionTree.csv', 'rU')
# 按行读取文件内容
reader = csv.reader(allElectronicsData)
# headers title一行
headers = reader.next()

print headers

# 特征值列表
featureList = []
# 标签列表
labelList = []

for row in reader:
	labelList.append(row[len(row)-1])
	# 每条样本特征值
	rowDict = {}
	for i in range(1, len(row)-1):
		# print row[i]
		rowDict[headers[i]] = row[i]
		# print "rowDict:", rowDict

	featureList.append(rowDict)

print featureList

# Vectorize features
vec = DictVectorizer()
# 转换特征值矩阵
dummyX = vec.fit_transform(featureList).toarray()

print "dummyX"+str(dummyX)
print vec.get_feature_names()

print "labelList:"+str(labelList)

# Vectorize class labels
lb = preprocessing.LabelBinarizer()
# 转换标签矩阵
dummyY = lb.fit_transform(labelList)
print "dummyY"+str(dummyY)

# Using decision tree for classifcation
# clf = tree.DecisionTreeClassifier()
clf = tree.DecisionTreeClassifier(criterion='entropy') # 信息熵（ID3）算法
# 训练决策树模型
clf = clf.fit(dummyX, dummyY)
print "clf:"+str(clf)

# Visulize model
# with open("allElectronicGiniOri.dot", 'w') as f:
with open("allElectronicInformationGainOri.dot", 'w') as f:
	# f = tree.export_graphviz(clf, out_file = f)
	# 绘制决策树模型 dot -Tpdf allElectronicInformationGainOri.dot -o DecisionTree.pdf
	f = tree.export_graphviz(clf, feature_names = vec.get_feature_names(), out_file=f)

# 第一个样本数据
oneRowX = dummyX[0, :]
print "oneRowX:"+str(oneRowX)

newRowX = oneRowX

# 修改数据称为新样本
newRowX[0] = 1
newRowX[2] = 0
print "newRowX"+str(newRowX)

# 决策树模型预测
predictedY = clf.predict(newRowX)
print "predictedY"+str(predictedY)